/*  电机编号
            头

     9  5       4 8
        1       0
          |dog|
          |dog|
          |dog|
        3       2
     11 7       6 10
*/

// double H_X = 0.715;  // (2.2)
// double H_Y = 0.25; // (2.2)
// double L_1 = 0.05;  // (2.2)
// double L_2 = 0.37;   // (2.2)
// double L_3 = 0.40;   // (2.2)

// typedef struct point
// {
//      float x;
//      float y;
//      float z;
// } bezier_Point;

#include <webots/Robot.hpp>
#include <webots/Motor.hpp>
#include <webots/PositionSensor.hpp>
#include <webots/InertialUnit.hpp>
#include <webots/Gyro.hpp>
#include <webots/Accelerometer.hpp>

#include <thread>
#include <atomic>
#include <chrono>

#include <cmath>
#include <iostream>
#include <string>
#include "global.h"
#include "all.h"
#include <qpOASES.hpp>
#include <Eigen/Dense>

#include <iomanip>

using namespace std;
using namespace Eigen;
using namespace webots;

// #define NUMBER_OF_LEDS 8
#define NUMBER_OF_JOINTS 12
// #define NUMBER_OF_CAMERAS 5

#define TIME_STEP 32
#define TROT_TEST
#define TEST_THREAD_STATE_ESTIMATE
// #define MPC_TEST

Matrix<double, 13 * PREDICTION_HORIZON, 1> *vecD;
Matrix<double, 13, 1> *x_0;
Matrix<double, 18, 1> *x_k;
Matrix<double, PREDICTION_HORIZON, 1> *psi_k_d;
Matrix<double, 3 * PREDICTION_HORIZON, 4> *r_k_ix;
Matrix<double, 12, 1> *f_MPC;
Matrix<double, 4, 3> *tau_motor;
Matrix<double, 4, 3> *q_motor;

int is_stand[4] = {1, 1, 1, 1};

static const char *motor_names[NUMBER_OF_JOINTS] = {
    "front right shoulder abduction motor", // 右前肩外展电机
    "front left shoulder abduction motor",  // 左前肩外展电机
    "rear right shoulder abduction motor",  // 右后肩外展电机
    "rear left shoulder abduction motor",   // 左后肩外展电机
    "front right shoulder rotation motor",  // 右前肩旋转电机
    "front left shoulder rotation motor",   // 左前肩旋转电机
    "rear right shoulder rotation motor",   // 右后肩旋转电机
    "rear left shoulder rotation motor",    // 左后肩旋转电机
    "front right elbow motor",              // 右前肘电机
    "front left elbow motor",               // 左前肘电机
    "rear right elbow motor",               // 右后肘电机
    "rear left elbow motor"                 // 左后肘电机
};

static Motor *motors[NUMBER_OF_JOINTS];
static PositionSensor *ps[NUMBER_OF_JOINTS];

void init_global_variables(Robot *robot) {
  // vecD.store(new Matrix<double, 13 * PREDICTION_HORIZON, 1>());
  // x_0.store(new Matrix<double, 13, 1>());
  // x_k.store(new Matrix<double, 18, 1>());
  // psi_k_d.store(new Matrix<double, PREDICTION_HORIZON, 1>());
  // r_k_ix.store(new Matrix<double, 3 * PREDICTION_HORIZON, 4>());
  vecD = new Matrix<double, 13 * PREDICTION_HORIZON, 1>();
  x_0 = new Matrix<double, 13, 1>();
  x_k = new Matrix<double, 18, 1>();
  psi_k_d = new Matrix<double, PREDICTION_HORIZON, 1>();
  r_k_ix = new Matrix<double, 3 * PREDICTION_HORIZON, 4>();
  f_MPC = new Matrix<double, 12, 1>();
  tau_motor = new Matrix<double, 4, 3>();
  q_motor = new Matrix<double, 4, 3>();

  for (int i = 0; i < NUMBER_OF_JOINTS; i++) {
    motors[i] = robot->getMotor(motor_names[i]);
    motors[i]->setPosition(0);
    ps[i] = motors[i]->getPositionSensor();
    ps[i]->enable(TIME_STEP);
  }
}

void setsmoothly(Robot *robot, const double *target, double duration) {
  double time_step = robot->getBasicTimeStep();
  const int n_steps_to_achieve_target =
      duration * 1000 / time_step; // 经过多少时间步完成动作

  double step_difference[NUMBER_OF_JOINTS];
  double current_position[NUMBER_OF_JOINTS];
  for (int i = 0; i < NUMBER_OF_JOINTS; ++i) {
    current_position[i] = motors[i]->getTargetPosition();
    step_difference[i] =
        (target[i] - current_position[i]) / n_steps_to_achieve_target;
  }
  for (int i = 0; i < n_steps_to_achieve_target; i++) {
    for (int j = 0; j < NUMBER_OF_JOINTS; ++j) { // 对每一个电机设置动作
      current_position[j] += step_difference[j];
      motors[j]->setPosition(current_position[j]);
    }
    // robot->step(time_step);
  }
}

Vector3d cubicBezier3D(Vector3d &p0, Vector3d &p1, Vector3d &p2, Vector3d &p3,
                       float t) {
  Vector3d result;
  for (int i = 0; i < 3; i++)
    result[i] = (1 - t) * (1 - t) * (1 - t) * p0[i] +
                3 * (1 - t) * (1 - t) * t * p1[i] +
                3 * (1 - t) * t * t * p2[i] + t * t * t * p3[i];
  return result;
}

// 足端轨迹曲线
Vector3d footTrajectory(Vector3d &start, Vector3d &end, double height,
                        double t_fai) {
  Vector3d result;
  result[0] =
      start[0] + (end[0] - start[0]) * (3 * pow(t_fai, 2) - 2 * pow(t_fai, 3));
  result[1] =
      start[1] + (end[1] - start[1]) * (3 * pow(t_fai, 2) - 2 * pow(t_fai, 3));

  double sigma_up = 2 * t_fai;       // 上升有效sigma
  double sigma_down = 2 * t_fai - 1; // 下降有效sigma
  double z_top = start[2] + height;  // 0.5时刻的高度

  if (sigma_up <= 1)
    result[2] =
        start[2] + height * (3 * pow(sigma_up, 2) - 2 * pow(sigma_up, 3));
  else if (sigma_down > 0)
    result[2] =
        z_top +
        (end[2] - z_top) * (3 * pow(sigma_down, 2) - 2 * pow(sigma_down, 3));

  return result;
}

// void reset_to_zero()
// {
//      cout << "reset_to_zero()" << endl;
//      Robot *robot = new Robot();

//      for (int i = 0; i < 12; i++)
//      {
//           Motor *motor = robot->getMotor(motor_names[i]);
//           motor->setPosition(0);
//           motor->getPositionSensor()->enable(TIME_STEP);
//      }
//      robot->step(TIME_STEP);

//      delete robot;
// }

void thread_trotForward(Robot *robot) {
  //  // p是起点，控制点1，控制点2，终点
  //  // y>0,L腿摆动相
  // Vector3d p0 = {-0.040, 0.060, -0.380};
  // Vector3d p1 = {0.102, 0.060, -0.234};
  // Vector3d p2 = {0.262, 0.060, -0.284};
  // Vector3d p3 = {0.306, 0.060, -0.380};

  // // L腿支撑相
  // Vector3d q0 = {0.306, 0.060, -0.380};
  // Vector3d q1 = {0.262, 0.060, -0.380};
  // Vector3d q2 = {0.122, 0.060, -0.380};
  // Vector3d q3 = {-0.040, 0.060, -0.380};

  Vector3d swing_start = {0.060, 0.060, -0.360};
  Vector3d swing_end = {0.266, 0.060, -0.360};
  double swing_height = 0.10;

  Vector3d point_st;
  Vector3d point_sw;

  Vector3d footposRF_0; // 0
  Vector3d footposLF_1; // 1
  Vector3d footposRR_2; // 2
  Vector3d footposLR_3; // 3

  int setStep = 1;
  while (robot->step(TIME_STEP) != -1) {
#ifdef TROT_TEST
    for (float t = 0.0; t <= 1.0; t += 0.1) {
      // cout<<"A"<<endl;
      // point_st = cubicBezier3D(p0, p1, p2, p3, t);
      // point_sw = cubicBezier3D(q0, q1, q2, q3, t);

      point_st = footTrajectory(swing_end, swing_start, 0, t);
      point_sw = footTrajectory(swing_start, swing_end, swing_height, t);

      footposRF_0 << point_st[0], -point_st[1], point_st[2];
      footposLR_3 << point_st[0] - 0.05, point_st[1], point_st[2];

      footposLF_1 << point_sw[0], point_sw[1], point_sw[2];
      footposRR_2 << point_sw[0] - 0.05, -point_sw[1], point_sw[2];

      double motors_target_pos[NUMBER_OF_JOINTS] = {
          inverseKinematics(0, footposRF_0)[0],
          inverseKinematics(1, footposLF_1)[0],
          inverseKinematics(2, footposRR_2)[0],
          inverseKinematics(3, footposLR_3)[0],
          inverseKinematics(0, footposRF_0)[1],
          inverseKinematics(1, footposLF_1)[1],
          inverseKinematics(2, footposRR_2)[1],
          inverseKinematics(3, footposLR_3)[1],
          inverseKinematics(0, footposRF_0)[2],
          inverseKinematics(1, footposLF_1)[2],
          inverseKinematics(2, footposRR_2)[2],
          inverseKinematics(3, footposLR_3)[2]};

      for (int i = 0; i < NUMBER_OF_JOINTS; i++) {
        motors[i]->setPosition(motors_target_pos[i]);
      }
      // setsmoothly(robot, motors_target_pos, 0.032 * setStep);
      robot->step(TIME_STEP);
    }

    for (float t = 0; t <= 1.0; t += 0.1) {
      // cout<<"B"<<endl;
      // point_st = cubicBezier3D(p0, p1, p2, p3, t);
      // point_sw = cubicBezier3D(q0, q1, q2, q3, t);

      point_st = footTrajectory(swing_end, swing_start, 0, t);
      point_sw = footTrajectory(swing_start, swing_end, swing_height, t);

      footposRF_0 << point_sw[0], -point_sw[1], point_sw[2];
      footposLR_3 << point_sw[0] - 0.05, point_sw[1], point_sw[2];

      footposLF_1 << point_st[0], point_st[1], point_st[2];
      footposRR_2 << point_st[0] - 0.05, -point_st[1], point_st[2];

      double motors_target_pos[NUMBER_OF_JOINTS] = {
          inverseKinematics(0, footposRF_0)[0],
          inverseKinematics(1, footposLF_1)[0],
          inverseKinematics(2, footposRR_2)[0],
          inverseKinematics(3, footposLR_3)[0],
          inverseKinematics(0, footposRF_0)[1],
          inverseKinematics(1, footposLF_1)[1],
          inverseKinematics(2, footposRR_2)[1],
          inverseKinematics(3, footposLR_3)[1],
          inverseKinematics(0, footposRF_0)[2],
          inverseKinematics(1, footposLF_1)[2],
          inverseKinematics(2, footposRR_2)[2],
          inverseKinematics(3, footposLR_3)[2]};

      for (int i = 0; i < NUMBER_OF_JOINTS; i++) {
        motors[i]->setPosition(motors_target_pos[i]);
      }
      // setsmoothly(robot, motors_target_pos, 0.032 * setStep);
      robot->step(TIME_STEP);
    }
#endif
    // motors[0]->setPosition(0.5);
    // motors[4]->setPosition(1.7);
    // motors[8]->setPosition(1.6);

    // Vector3d footPosition(0.346, -0.05, -0.40);
    // Vector3d motorAngles = inverseKinematics(0, footPosition);
    // cout<<motorAngles<<endl;
    // motors[0]->setPosition(motorAngles[0]);
    // motors[4]->setPosition(motorAngles[1]);
    // motors[8]->setPosition(motorAngles[2]);
  }
}

void thread_state_estimate(Robot *robot, StateEstimator *se, GaitPlanner *gp) {

  cout << "Start thread_state_estimate." << endl;
  // 初始化传感器(在init_global_variables()中)

  InertialUnit *iu = robot->getInertialUnit("inertial unit");
  iu->enable(TIME_STEP);
  Gyro *gyro = robot->getGyro("gyro");
  gyro->enable(TIME_STEP);
  Accelerometer *acc = robot->getAccelerometer("accelerometer");
  acc->enable(TIME_STEP);

  cout << "Initialization finished." << endl;

  // 获取IMU数据
  Vector3d Eulerangles;
  Vector3d a_original;
  Vector3d omega_original;

  // 编码器获取电机位置和速度
  Matrix<double, 4, 3> q_j_i;
  Matrix<double, 4, 3> q_dot_j_i;

  Matrix<double, 4, 3> q_j_i_last;
  q_j_i_last.setZero();

  while (robot->step(TIME_STEP) != -1) {
    // for (int i = 0; i < 4; i++) {
    //   q_j_i(i, 0) = ps[3 * i]->getValue();
    //   q_j_i(i, 1) = ps[3 * i + 1]->getValue();
    //   q_j_i(i, 2) = ps[3 * i + 2]->getValue();
    //   q_dot_j_i(i, 0) = motors[3 * i]->getVelocity();
    //   q_dot_j_i(i, 1) = motors[3 * i + 1]->getVelocity();
    //   q_dot_j_i(i, 2) = motors[3 * i + 2]->getVelocity();
    // }

    Eulerangles = Map<Vector3d>(const_cast<double *>(
        iu->getRollPitchYaw())); // BUG: IMU零系不是通常理解的z轴为偏航
    a_original = Map<Vector3d>(const_cast<double *>(acc->getValues()));
    omega_original = Map<Vector3d>(const_cast<double *>(gyro->getValues()));

    // cout << "IMU data got." << endl;

    for (int i = 0; i < 3; i++) {
      q_j_i(0, i) = ps[4 * i]->getValue();
      q_j_i(1, i) = ps[4 * i + 1]->getValue();
      q_j_i(2, i) = ps[4 * i + 2]->getValue();
      q_j_i(3, i) = ps[4 * i + 3]->getValue();

      // q_dot_j_i(0, i) = motors[4 * i]->getVelocity();
      // q_dot_j_i(1, i) = motors[4 * i + 1]->getVelocity();
      // q_dot_j_i(2, i) = motors[4 * i + 2]->getVelocity();
      // q_dot_j_i(3, i) = motors[4 * i + 3]->getVelocity(); // BUG:
      // webots里没有速度传感器，这个不是实时电机转速

      q_dot_j_i(0, i) = (q_j_i(0, i) - q_j_i_last(0, i)) / 0.064;
      q_dot_j_i(1, i) = (q_j_i(1, i) - q_j_i_last(1, i)) / 0.064;
      q_dot_j_i(2, i) = (q_j_i(2, i) - q_j_i_last(2, i)) / 0.064;
      q_dot_j_i(3, i) = (q_j_i(3, i) - q_j_i_last(3, i)) / 0.064;

      q_j_i_last(0, i) = q_j_i(0, i);
      q_j_i_last(1, i) = q_j_i(1, i);
      q_j_i_last(2, i) = q_j_i(2, i);
      q_j_i_last(3, i) = q_j_i(3, i);
    }

    cout << "q_j_i: " << q_j_i << endl << endl;
    cout << "q_dot_j_i: " << q_dot_j_i << endl << endl;

    // cout << "Encoder data got." << endl;

    se->setIMUdata(Eulerangles, a_original, omega_original);
    se->setEncoderdata(q_j_i, q_dot_j_i);
    se->calculateAll();
    se->kalman_update(10);

    cout << "State estimation finished." << endl;
    // 获取状态估计值
    *x_k = se->getx_k();

    x_0->segment(0, 3) = Eulerangles;
    x_0->segment(3, 3) = x_k->segment(0, 3);
    x_0->segment(6, 3) = omega_original;
    x_0->segment(9, 3) = x_k->segment(3, 3);
    (*x_0)(12) = -9.8;

#ifdef TEST_THREAD_STATE_ESTIMATE
    // cout<<"x_0: "<<x_0->transpose()<<endl;
    cout << "------------------x_0-----------------:" << endl;
    cout << "位置:" << x_0->segment(3, 3).transpose() << endl;
    cout << "速度:" << x_0->segment(9, 3).transpose() << endl;
    cout << "角度: " << x_0->segment(0, 3).transpose() << endl;
    cout << "角速度: " << x_0->segment(6, 3).transpose() << endl;
#endif

    // for (int i = 0; i < 4; i++) {
    //   tau_motor->row(i) =
    //       (f_MPC->segment(i, 3)).transpose() * JacobianMatrix(q_j_i.row(i),
    //       i);
    // }
  }
}

int main() {
  // Vector3d footPosition(-0.25, 0.05, -0.40);
  // cout<<inverseKinematics(1, footPosition)<<endl;
  // // cout<<Antisymmetric(Vector3d(1, 2, 3))<<endl;
  // cout<<footPosition<<endl;
  // reset_to_zero();

  Robot *robot = new Robot();
  GaitPlanner *gp = new GaitPlanner();
  // TrajectoryGenerator *tg = new TrajectoryGenerator();
  StateEstimator *se = new StateEstimator();
  // MPCsolver *mpc = new MPCsolver();

  init_global_variables(robot);
  thread thread1(thread_trotForward, robot);
  thread thread2(thread_state_estimate, robot, se, gp);

  thread1.join();
  thread2.join();

  // delete robot;
  delete gp;
  // delete tg;
  delete se;
  // delete mpc;
}